from pathlib import Path
import loguru

class DecrypteFile:
    @staticmethod
    def decrypted_vmd(input_vmd_file_path: Path) -> Path:
        loguru.logger.info(f"解密文件：{input_vmd_file_path}")
        input_path = Path(input_vmd_file_path)

        loguru.logger.debug('正在验证输入文件中……')
        DecrypteFile._validate_input_file(input_path)

        loguru.logger.debug('生成输出文件路径……')
        output_path = DecrypteFile._generate_output_path(input_path)

        known_plaintexts = [
            b'Vocaloid Motion Data 0002',
            b'Vocaloid Motion Data file'
        ]
        loguru.logger.debug('检测密钥中……')
        key = DecrypteFile._detect_key(input_path, known_plaintexts)

        loguru.logger.debug('解密文件中……')
        DecrypteFile._decrypt_file(input_path, output_path, key)

        loguru.logger.success(f"解密完成，输出文件：{output_path}")
        loguru.logger.success('请注意：解密后的文件可能无法正常使用，需要手动检查和修复。')
        return output_path

    @staticmethod
    def decrypted_pmx(input_pmx_file_path: Path) -> Path:
        loguru.logger.info(f"解密文件：{input_pmx_file_path}")
        input_path = Path(input_pmx_file_path)

        loguru.logger.debug('正在验证输入文件中……')
        DecrypteFile._validate_input_file(input_path)

        loguru.logger.debug('生成输出文件路径……')
        output_path = DecrypteFile._generate_output_path(input_path)

        known_plaintexts = [
            b'PMX',
        ]
        loguru.logger.debug('检测密钥中……')
        key = DecrypteFile._detect_key(input_path, known_plaintexts)

        loguru.logger.debug('解密文件中……')
        DecrypteFile._decrypt_file(input_path, output_path, key)

        loguru.logger.success(f"解密完成，输出文件：{output_path}")
        loguru.logger.success('请注意：解密后的文件可能无法正常使用，需要手动检查和修复。')
        return output_path

    @staticmethod
    def _generate_output_path(input_path: Path) -> Path:
        output_filename = input_path.stem + '_out' + input_path.suffix
        return input_path.with_name(output_filename)

    @staticmethod
    def _validate_input_file(input_path: Path):
        if not input_path.is_file():
            raise FileNotFoundError(f"输入文件不存在：{input_path}")

    @staticmethod
    def _detect_key(input_path: Path, known_plaintexts: list) -> int:
        max_known_length = max(len(kp) for kp in known_plaintexts)
        with input_path.open('rb') as f_in:
            encrypted_header = f_in.read(max_known_length)

            for known_plaintext in known_plaintexts:
                known_length = len(known_plaintext)
                encrypted_segment = encrypted_header[:known_length]

                possible_keys = [eb ^ kb for eb, kb in zip(encrypted_segment, known_plaintext)]
                if len(set(possible_keys)) == 1:
                    detected_key = possible_keys[0]
                    print(f"使用已知明文 '{known_plaintext.decode('ascii')}' 检测到的密钥：0x{detected_key:02X}")
                    return detected_key

        raise ValueError("无法检测到唯一的密钥，可能需要更新已知的明文列表。")

    @staticmethod
    def _decrypt_file(input_path: Path, output_path: Path, key: int):
        with input_path.open('rb') as f_in, output_path.open('wb') as f_out:
            while True:
                chunk = f_in.read(4096)
                if not chunk:
                    break
                decrypted_chunk = bytes([b ^ key for b in chunk])
                f_out.write(decrypted_chunk)


if __name__ == '__main__':
    # vmd_file_path = Path('test.vmd')
    # decrypted_vmd_file_path = DecrypteFile.decrypted_vmd(vmd_file_path)
    # print(f"解密后的文件：{decrypted_vmd_file_path}")

    pmx_file_path = Path(
        r"D:\steam\steamapps\workshop\content\1968650\2803712418\TDA 幻想金瓶梅[02] 瓶兒 Ver1.00[Silver].pmx")
    decrypted_pmx_file_path = DecrypteFile.decrypted_pmx(pmx_file_path)
    print(f"解密后的文件：{decrypted_pmx_file_path}")
